import os
import numpy as np
import tensorflow as tf
from PIL import Image

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']
CHAR_SET = number + alphabet + ALPHABET
CHAR_SET_LEN = len(CHAR_SET)
MAX_CAPTCHA = 4
IMAGE_HEIGHT = 26
IMAGE_WIDTH= 85
SAVE_PATH = "D:\\PythonProject\\vc_recognition_test\\"
BASE_PATH = "D:\\PythonProject\\vc_recognition_test\\vc_base1\\"
#SUM = len(all_image = os.listdir("D:\\PythonProject\\vc_recognition_test\\vc_base1\\"))

def get_name_and_Image(index):
    all_image = os.listdir(BASE_PATH)
    text = os.path.splitext(all_image[index])[0]
    image = Image.open(BASE_PATH + all_image[index])
    image = np.array(image)
    return text,image

def text2vec(text):
    vector = np.zeros([MAX_CAPTCHA,CHAR_SET_LEN])
    for i, c in enumerate(text):
        idx = CHAR_SET.index(c)
        vector[i][idx] = 1.0
    return vector

def vec2text(vec):
    text = []
    for i, c in enumerate(vec):
        text.append(CHAR_SET[c])
    return "".join(text)

def get_next_batch(index):
    batch_x = np.zeros([1, IMAGE_HEIGHT,IMAGE_WIDTH,1])
    batch_y = np.zeros([1, MAX_CAPTCHA,CHAR_SET_LEN])
    text,image = get_name_and_Image(index)
    image = np.mean(image, -1)
    batch_x[0, : , : ,0] = image
    batch_y[0, : , : ] = text2vec(text)
    return batch_x,batch_y

def predict():
    model = tf.keras.models.load_model(SAVE_PATH + 'model')
    success = 0
    count = 100
    for i in range(count):
        data_x, data_y = get_next_batch(i)
        prediction_value = model.predict(data_x)
        data_y = vec2text(np.argmax(data_y, axis=2)[0])
        prediction_value = vec2text(np.argmax(prediction_value, axis=2)[0])
        if data_y.upper() == prediction_value.upper():
            print("预测=", prediction_value, "实际=", data_y, "预测成功。")
            success += 1
        else:
            print("预测=", prediction_value, "实际=", data_y, "预测失败。")
    print("预测", count, "次", "成功", success, "次", "成功率 =", success / count)

def main():
    predict()

if __name__ == '__main__':
    main()
